import cv2  # OpenCV for video processing
import os   # For handling file paths and directories
import glob  # For finding files with specific extensions

def extract_first_frame(video_path, output_dir):
    """
    Extract the first frame from a video and save it as an image.
    
    :param video_path: Path to the video file.
    :param output_dir: Directory to save the output image.
    """
    # Open the video file
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        print(f"Error: Could not open video file {video_path}")
        return  # Exit if the video can't be opened
    
    # Read the first frame
    ret, frame = cap.read()  # ret is a boolean indicating success
    
    if ret:
        # Get the base name of the video file and create a new output file name
        base_name = os.path.basename(video_path)  # e.g., "video1.mp4"
        name, ext = os.path.splitext(base_name)  # Split into name and extension
        output_path = os.path.join(output_dir, f"{name}_frame1.jpg")  # e.g., "video1_frame1.jpg"
        
        # Save the frame as an image
        cv2.imwrite(output_path, frame)
        print(f"Successfully saved: {output_path}")
    else:
        print(f"No frames found in video: {video_path}")
    
    # Release the video capture object
    cap.release()

def main():
    """
    Main function to prompt user for inputs and extract frames.
    """
    # Prompt user for input directories
    video_dir = input("请输入视频目录路径（例如，/Users/yourname/videos）：")
    output_dir = input("请输入输出目录路径（例如，/Users/yourname/output_images）：")
    
    # Create the output directory if it doesn't exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created output directory: {output_dir}")
    
    # Define common video extensions
    video_extensions = ['*.mp4', '*.avi', '*.mov', '*.mkv']  # Add more if needed
    
    # Get a list of all video files in the directory
    video_files = []
    for ext in video_extensions:
        video_files.extend(glob.glob(os.path.join(video_dir, ext)))  # glob.glob finds files matching the pattern
    
    if not video_files:
        print(f"No video files found in {video_dir}. Please check the directory.")
        return  # Exit if no videos are found
    
    # Process each video file
    for video in video_files:
        extract_first_frame(video, output_dir)

if __name__ == "__main__":
    main()
